using Oceananigans
using NumericalEarth.DataWrangling: NearestNeighborInpainting, inpaint_mask!

grid = RectilinearGrid(CPU(); size=(10, 10, 1), x=(0, 1), y=(0, 1), z=(0, 1))

# --- Bug 1: value == 0 vs donors == 0 ---
# A masked cell surrounded by neighbors that average to zero
# (e.g. temperature data near 0°C with values -2, +2, -2, +2)
field = CenterField(grid)
mask = Field{Center, Center, Center}(grid, Bool)
interior(field) .= 0.0
interior(field)[4, 5, 1] = -2.0
interior(field)[6, 5, 1] =  2.0
interior(field)[5, 4, 1] = -2.0
interior(field)[5, 6, 1] =  2.0
interior(mask)[5, 5, 1] = true

inpaint_mask!(field, mask; inpainting=NearestNeighborInpainting(100))
@show interior(field)[5, 5, 1]
# Before fix: NaN → then 0 (via _fill_nans!). Wrong code path.
# After fix:  0.0 directly from averaging. Correct.

# --- Bug 2: _fill_nans! converts NaN → 0 instead of field mean ---
# A salinity-like field (~34 psu) with a large masked region
# and finite maxiter so the center can't be reached by propagation
field2 = CenterField(grid)
mask2 = Field{Center, Center, Center}(grid, Bool)
set!(field2, (x, y, z) -> 34.0)
for i in 3:7, j in 3:7
    interior(mask2)[i, j, 1] = true
end
inpaint_mask!(field2, mask2; inpainting=NearestNeighborInpainting(1))
@show interior(field2)[5, 5, 1]
# Before fix: 0.0 (from _fill_nans! multiplying by !isnan). Catastrophic for salinity.
# After fix:  34.0 (filled with field mean). Physically reasonable.